'''Functions to put dust events into hdf5 files.

Mike Dorey  2011-03-24
'''


import numpy
import h5py
import datetime


STR_DTYPE = h5py.new_vlen(str)
DUST_EVENT_DTYPE = numpy.dtype([('timestamp', '<u8'),
                                ('velocity', '<f'),
                                ('charge', '<f'),
                                ('mass', '<f'),
                                ('terminalpotential', '<f'),
                                ('lenspotential', '<f'),
                                ('needlepotential', '<f'),
                                ('pulsefreq', '<f'),
                                ('pulsewidth', '<f'),
                                ('pulseheight', '<f'),
                                ('selection', '<u1'),
                                ('waveforms', STR_DTYPE)])


def open_hdf5_file(hdf5_filename, file_version):
    '''Opens a HDF5 file and returns the file pointer. It
    also sets the attributes of the '/' group.

    Arguments
        hdf5_filename - name of the output hdf5 file
        file_version - version of the file
    '''
    f = h5py.File(hdf5_filename, 'w')
    f['/'].attrs['file version'] = file_version
    f['/'].attrs['when created'] = get_time_now()
    return f


def create_user_group(hdf5_file, user_name):
    '''Creates a group in the HDF5 file with the user's name.
    Returns the group

    Arguments
        user_name - name of the user
    '''
    user_grp = hdf5_file.create_group(user_name)
    return user_grp


def add_experiments_to_user_group(user_group, experiments):
    '''Adds experiments to the user's HDF5 group.

    Arguments
        user_group - HDF5 group for the user
        experiments - list of DustExperiment objects
    '''
    for experiment in experiments:
        create_experiment_group(user_group, experiment) 


def get_time_now():
    '''Returns a string of the form "mm/dd/yyyy hh:mm:ss" for the
    current time. This is UTC time.
    '''
    now = datetime.datetime.utcnow()
    return now.strftime('%02m/%02d/%4Y %02H:%02M:%02S')


def create_experiment_group(user_grp, experiment):
    '''Creates an experiment group in the hdf5 file.

    Arguments
        user_grp - the hdf5 user group
        experiment - a DustExperiment object
    '''
    np_events = to_numpy_records(experiment.dust_events)
    ds = user_grp.create_dataset(experiment.experiment_name, data=np_events, 
                   dtype=DUST_EVENT_DTYPE)
    ds.attrs['remarks'] = experiment.remarks
    ds.attrs['chemical comp'] = experiment.dust_type.chemical_comp
    ds.attrs['batch'] = experiment.dust_type.batch


def create_dust_events_dataset(hdf5_file, config, dust_events):
    '''Creates a dust events dataset in the root group of the
    HDF5 file.

    Arguments
        hdf5_file - HDF5 file.
        config - Dictionary that holds the configuration parameters.
            These will be put into attributes of the dataset.
        dust_events - List of DustEvent objects that will be put
            into the dataset.
    '''
    np_events = to_numpy_records(dust_events)
    ds = hdf5_file.create_dataset('Dust events', data=np_events, 
                   dtype=DUST_EVENT_DTYPE)
    for key in config:
        ds.attrs[key] = config[key]


def to_numpy_records(events):
    '''Returns a numpy array of records that can be used in creating
    an hdf5 dataset.

    Argument
        events - list of DustEvent objects
    '''
    records = numpy.zeros((len(events),), dtype=DUST_EVENT_DTYPE)
    for (i, e) in enumerate(events):
        records[i]['timestamp'] = e.timestamp
        records[i]['velocity'] = e.velocity
        records[i]['charge'] = e.charge
        records[i]['mass'] = e.mass
        records[i]['terminalpotential'] = e.terminalpotential
        records[i]['lenspotential'] = e.lenspotential
        records[i]['needlepotential'] = e.needlepotential
        records[i]['pulsefreq'] = e.pulsefreq
        records[i]['pulsewidth'] = e.pulsewidth
        records[i]['pulseheight'] = e.pulseheight
        records[i]['selection'] = e.selection
        records[i]['waveforms'] = e.waveforms
    return records
